import numpy as np
from skimage.io import imsave, imread
from skimage import img_as_float, img_as_ubyte


# --------- 文件处理和二进制转换 --------- #

# 输入：文件名
# 输出：作为字符串的分组位列表
def file_to_bin(filename, bits=1):
    bin_file = open(filename, 'rb').read()
    print(bin_file)
    bin_strings = map(bin, list(bin_file))

    binlist = list(''.join([i[2:].rjust(8, "0") for i in bin_strings]))

    if bits == 1:
        return binlist

    answer, s = list(), ''
    for i in range(len(binlist)):
        s += binlist[i]
        if i % bits == bits - 1:
            answer.append(s)
            s = ''
    return answer

def mes_to_bin(mes, bits=1):
    bin_mes = bytes(mes, encoding='utf-8')
    print(bin_mes)
    print(len(bin_mes))
    if len(bin_mes) > 256 ** 2:
        print("隐写消息太大")
        return []
    else:
        a = bin(len(bin_mes) // 256)
        b = bin(len(bin_mes) % 256)
        c = list(a[2:].rjust(8, "0"))
        d = list(b[2:].rjust(8, "0"))

    bin_strings = map(bin, list(bin_mes))

    binlist = c + d + c + d + list(''.join([i[2:].rjust(8, "0") for i in bin_strings]))

    if bits == 1:
        return binlist

    answer, s = list(), ''
    for i in range(len(binlist)):
        s += binlist[i]
        if i % bits == bits - 1:
            answer.append(s)
            s = ''
    return answer

def bin_to_mes(binlist):
    s = ''
    bytes_list = list()
    for i in range(len(binlist)):
        s += binlist[i]
        if i % 8 == 7:
            bytes_list.append(int(s, base=2))
            s = ''
    return map(str.encode(), bytes_list)

# Input:  list of grouped bits as strings
# Output: the file with the specified name is saved
def bin_to_file(binlist, filename='output_file.txt'):
    s = ''
    bytes_list = list()
    a = b = 0
    for i in range(16):
        s += binlist[i]
        if i % 8 == 7:
            a = int(s, base=2) + a * 256
            s = ''

    for i in range(16, 32):
        s += binlist[i]
        if i % 8 == 7:
            b = int(s, base=2) + b * 256
            s = ''
    if a != b:
        return 'error_of_bits!'
    for i in range(32, a * 8 + 32):
        s += binlist[i]
        if i % 8 == 7:
            bytes_list.append(int(s, base=2))
            s = ''

    binary_data = bytes(bytes_list)
    data = str(binary_data, encoding="utf-8")
    out_file = open(filename, 'wb')
    out_file.write(binary_data)
    out_file.close()
    return data


# ------------------- PVD功能-------------------- #

# 输入：两个像素之间的差值
# 输出：要嵌入的位数，左侧和右侧
# 差异所在区间的边界
def embed_number(n, i=0):
    # srange = (0, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 192, 256)
    srange = [[0, 8, 16, 32, 64, 128, 256]
        , [0, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 172, 256]
        , [0, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 172, 256]]

    l, r = 0, len(srange[i]) - 1
    while r - l > 1:
        mid = (l + r) // 2
        if srange[i][mid] >= n:
            r = mid
        else:
            l = mid
    return int(np.log2(srange[i][r] - srange[i][l])), srange[i][l], srange[i][r]


# 输入：[0；255]中的三个数字
# 输出：[0；255]中的两个数字被改变
def change_difference(a, b, dif, newdif):
    swap = False
    if a > b:
        a, b = b, a
        swap = True
    upper_add = abs(newdif - dif) // 2 + abs(newdif - dif) % 2
    lower_add = abs(newdif - dif) // 2
    # 越界处理
    if newdif > dif:
        if a - upper_add < 0:
            shift = upper_add - a
            upper_add -= shift
            lower_add += shift
        if b + lower_add > 255:
            shift = b + lower_add - 255
            lower_add -= shift
            upper_add += shift
        a -= upper_add
        b += lower_add
    else:
        a += upper_add
        b -= lower_add
    if swap:
        a, b = b, a
    return a, b


def pixel_dif(a, b):
    return max(a, b) - min(a, b)


# -------------- 额外功能 -------------- #

# 输入：两个相同形状的图像
# 输出：差分图像（以及每个通道的三个差分图像，如果存在）
def generate_difference(original, output):
    # 找出频道数
    c = 3 if len(original.shape) == 3 else 1
    dif = abs(img_as_float(original) - img_as_float(output))
    if c == 1:
        dif -= dif.min()
        dif = dif / dif.max()
        imsave("difference.png", img_as_ubyte(dif))
    else:
        black = original[:, :, 0] * 0
        dif_r = dif[:, :, 0]
        dif_g = dif[:, :, 1]
        dif_b = dif[:, :, 2]
        dif = np.dstack((dif_b, dif_g, dif_r))
        dif -= dif.min()
        dif = dif * 1 / dif.max()
        dif_r -= dif_r.min()
        dif_r = dif_r / dif_r.max()
        dif_r = np.dstack((dif_r, black, black))
        dif_g -= dif_g.min()
        dif_g = dif_g / dif_g.max()
        dif_g = np.dstack((black, dif_g, black))
        dif_b -= dif_b.min()
        dif_b = dif_b / dif_b.max()
        dif_b = np.dstack((black, black, dif_b))
        imsave("difference.png", img_as_ubyte(dif))
        imsave("difference_r.png", img_as_ubyte(dif_r))
        imsave("difference_g.png", img_as_ubyte(dif_g))
        imsave("difference_b.png", img_as_ubyte(dif_b))
    return img_as_ubyte(dif)

if __name__ == '__main__':


    print(mes_to_bin("Hello,World!中", 2))
